import sys
import argparse
import os

# https://www.codewithgpu.com/u/iscyy

# sys.path.append('/root/ultralyticsPro/') # Path 以Autodl为例

from ultralytics import YOLO

def main(opt):
    yaml = opt.cfg 
    # model = YOLO(yaml) # 直接加载yaml文件训练
    epochs = opt.epochs
    # model = YOLO(weights)  # 直接加载权重文件进行训练
    model = YOLO(yaml).load('yolo11n.pt') # 加载yaml配置文件的同时，加载权重进行训练

    model.info()
# /kaggle/input/facere2/face_reduced/data.yaml
    results = model.train(data='/kaggle/working/v11/expression.yaml',  # 训练参数均可以重新设置
                        epochs=epochs, 
                        imgsz=640, 
                        # workers=8, 
                        cache=True,
                        # batch=2,
                        ## 保存路径
                        project='runs/train/'+yaml.split('/')[-1].split('.')[0],
                        ## 保存间隔
                        save_period=5,
                        ## 是否使用数据增强
                        # augment=True,
                        # verbose=True,
                        # val=True,
                        # plots=True,
                        )

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='ultralytics/cfg/models/11/yolo11.yaml', help='initial weights path')
    parser.add_argument('--weights', type=str, default='', help='')
    parser.add_argument('--epochs', type=int, default=10, help='')

    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt

if __name__ == "__main__":
    opt = parse_opt()
    main(opt)